#include <unittest/unittest.h>

#include <cusp/array2d.h>
#include <cusp/multiply.h>
#include <cusp/device/block/qr.h>


template<typename MatrixType1, typename MatrixType2, typename MatrixType3>
void checkQRresult(MatrixType1 & A, MatrixType2 & Q, MatrixType3 & R)
{
	/* check that R is upper triangular */
	for (size_t i = 0; i < R.num_rows; ++i)
		for (size_t j = 0; j < i; ++j)
			ASSERT_ALMOST_EQUAL(R(i,j), 0.0);
	
	
	/* check decomposition */
	cusp::array2d<float, cusp::host_memory> Ch;
	cusp::multiply(Q, R, Ch);
	ASSERT_EQUAL(Ch.num_rows, A.num_rows);
	ASSERT_EQUAL(Ch.num_cols, A.num_cols);
	for (size_t i = 0; i < Ch.num_rows; ++i)
		for (size_t j = 0; j < Ch.num_cols; ++j)
			ASSERT_ALMOST_EQUAL(A(i,j), Ch(i,j));
}

template<typename RandomAccessIterator1, typename RandomAccessIterator2>
__global__
void testQR_kernel(RandomAccessIterator1 A, size_t num_rows, size_t num_cols, RandomAccessIterator2 Q)
{
	if (blockIdx.x == 0)
		cusp::device::block::qr(A, num_rows, num_cols, Q);
}

void TestQRDeviceBlockHelper(cusp::array2d<float, cusp::host_memory> & Ah)
{
	cusp::array2d<float, cusp::device_memory> A = Ah;
	cusp::array2d<float, cusp::device_memory, cusp::column_major> Q(A.num_rows, A.num_rows, 0.0);
	testQR_kernel<<<1,128>>>(A.values.begin(), A.num_rows, A.num_cols, Q.values.begin());
	cusp::array2d<float, cusp::host_memory> Rh = A;
	cusp::array2d<float, cusp::host_memory> Qh = Q;
	checkQRresult(Ah, Qh, Rh);
}

void TestQRDeviceBlock(void)
{
	/* test with more rows than columns */

	cusp::array2d<float, cusp::host_memory> Ah(5,3);
	Ah(0,0) = 12.0; Ah(0,1) = -51.0;  Ah(0,2) =   4.0;
	Ah(1,0) =  6.0; Ah(1,1) = 167.0;  Ah(1,2) = -68.0;
	Ah(2,0) = -4.0; Ah(2,1) =  24.0;  Ah(2,2) = -41.0;
	Ah(3,0) =  8.0; Ah(3,1) =  45.0;  Ah(3,2) =   0.0;
	Ah(4,0) = -4.0; Ah(4,1) = -86.0;  Ah(4,2) =   1.0;

	TestQRDeviceBlockHelper(Ah);
}
DECLARE_UNITTEST(TestQRDeviceBlock);

void TestQRDeviceBlock2(void)
{
	/* test with more columns than rows */

	cusp::array2d<float, cusp::host_memory> Ah(3,5);
	Ah(0,0) = 12.0; Ah(1,0) = -51.0;  Ah(2,0) =   4.0;
	Ah(0,1) =  6.0; Ah(1,1) = 167.0;  Ah(2,1) = -68.0;
	Ah(0,2) = -4.0; Ah(1,2) =  24.0;  Ah(2,2) = -41.0;
	Ah(0,3) =  8.0; Ah(1,3) =  45.0;  Ah(2,3) =   0.0;
	Ah(0,4) = -4.0; Ah(1,4) = -86.0;  Ah(2,4) =   1.0;

	TestQRDeviceBlockHelper(Ah);
}
DECLARE_UNITTEST(TestQRDeviceBlock2);

